/*
 *  ballFinder.cpp
 *  Ball and Plate
 *
 *  Created by Kenn Sebesta on 2007-11-23.
 *  Copyright 2007 __MyCompanyName__. All rights reserved.
 *
 */

#include "ballFinder.h"
#include <OpenCV/OpenCV.h>
#include <iostream>

CvMat *convMat; //This was declared global because I couldn't figure out any other way to kill the memory leak.

void PrintMat (CvMat *);

void ballFinder(/*Input variables*/		IplImage *maskImage, IplImage *camFrame, CvRect cropRect, int threshLow, int threshHigh,
					 /*Output variables*/	double *convMinOUT, double *convMaxOUT, CvMat *y_measrd, CvMat **convMatOUT, IplImage **edgeImageOUT, IplImage **bwImageOUT)
{
	// Intialize image matrices
	IplImage *cropImage, *grayImage, *bwImage, *edgeImage;
	if( cropRect.width!=0 )
		cvSetImageROI( camFrame, cropRect ); //crops image.
	else
		cvSetImageROI( camFrame, cvRect( 0, 0, camFrame->width, camFrame->height ));

	cropImage = cvCreateImage(cvGetSize(camFrame), camFrame->depth, camFrame->nChannels);	
	grayImage = cvCreateImage(cvGetSize(cropImage), cropImage->depth, 1);
	bwImage = cvCreateImage(cvGetSize(cropImage), cropImage->depth, 1);
	edgeImage = cvCreateImage(cvGetSize(cropImage), cropImage->depth, 1);	
	
	cvCopy( camFrame, cropImage );
	if (cropImage->nChannels!=1){
		cvSplit( cropImage, grayImage, NULL, NULL, NULL );
		//	cvCvtColor(cropImage, grayImage, CV_BGR2GRAY);	//Can't figure out how to tell if it's a RGB or 
																			//grayscale because my camera returns three 
																			//channels, even though it's grayscale.		
	}
	else{
		cvCopy(cropImage,grayImage);
	}

	
	//Image analysis
/*	if(0)
		cvThreshold(grayImage, bwImage, valueThr1, valueThr2, CV_THRESH_BINARY); // Binarises the image. HAVEN'T IMPLEMENTED YET.
	else 
		cvCopy( grayImage, bwImage );*/
	cvCopy( grayImage, bwImage );
	cvCanny( bwImage, edgeImage, threshLow, threshHigh, 3 ); // Performs the edge extraction on the image
	
	int dft_M = cvGetOptimalDFTSize( edgeImage->height + maskImage->height - 1 );	//This code could be done offline, but there's no reason to because it's so fast.
	int dft_N = cvGetOptimalDFTSize( edgeImage->width + maskImage->width - 1 );		//This code could be done offline, but there's no reason to because it's so fast.
	conv2D( edgeImage, maskImage, dft_M, dft_N ); // Performs the convolution of the images. The convolution matrix is convMat, which is global.

	
	//================================================
	// Returns the maximum convolution point, along with the max/min convolution values
	// in order to properly scale the output image.		
	//================================================
//	double minConv;
//	CvPoint *ballPoint; 
//		ballPoint=(CvPoint *) malloc(sizeof(CvPoint));	
//	cvMinMaxLoc( convMat, &minConv, convMaxOUT, NULL, ballPoint);

	CvPoint ballPoint; //Memory Leak? Don't know how to release this from memory.
	cvMinMaxLoc( convMat, convMinOUT, convMaxOUT, NULL, &ballPoint);

/*	if ( *convMaxOUT == 3.5e06 ) //This needs to be modularized. Right now, this value needs to be changed here AND in run.cpp
		return; */	// CANNOT DO THIS HERE BECAUSE WE DON'T FREE UP THE MEMORY. NEED SOME SORT OF goto.
						// MIGHT BE OTHER ISSUES, TOO. NEED TO INVESTIGATE.

	//================================================
	// Perform barycenter
	//================================================
	
	int pixelNeighborhood=5, x0, x1, y0, y1;
	double weight;
	CvMat *xLeverArm, *yLeverArm, *xSumTmp, *ySumTmp;
	CvMat baryMat;
	
// Calculation of a small window around ballPoint, with a verification that we don't leave the convolution image.
	x0=std::max( ballPoint.x - pixelNeighborhood, 0 );
	x1=std::min( ballPoint.x + pixelNeighborhood, edgeImage->width-1);
	y0=std::max( ballPoint.y - pixelNeighborhood, 0 );
	y1=std::min( ballPoint.y + pixelNeighborhood, edgeImage->height-1);
	xLeverArm=cvCreateMat( x1-x0+1, 1, CV_32F );
	yLeverArm=cvCreateMat( 1, y1-y0+1, CV_32F );
	xSumTmp=cvCreateMat( y1-y0+1, 1, CV_32F );
	ySumTmp=cvCreateMat( 1, x1-x0+1, CV_32F );
	for ( int i=x0; i<=x1 ; i++ ) 
		cvmSet( xLeverArm, i-x0, 0, (double)i ); //cvmSet is a tricky function. It references from 0.
	for ( int i=y0; i<=y1 ; i++ ) 
		cvmSet( yLeverArm, 0, i-y0, (double)i ); //cvmSet is a tricky function. It references from 0.

//-----------------------------
//	CvFileStorage* fs = cvOpenFileStorage( "crash.yml", 0, CV_STORAGE_WRITE);
//	cvWrite( fs, "xLeverArm", xLeverArm );
//	cvWrite( fs, "yLeverArm", yLeverArm );
//-----------------------------
	
//	printf("pixel neighborhood is x=%d, y=%d, width=%d, height=%d.\n", x0, y0, x1-x0+1, y1-y0+1);
//	printf("convMat size is width=%d, height=%d. The \"ball\" is at a relative (%d,%d)\n", convMat->rows, convMat->cols, ballPoint.x, ballPoint.y);
	
	cvGetSubRect( convMat, &baryMat, cvRect( x0, y0, x1-x0+1, y1-y0+1 ) ); //crops image

//	cvWrite( fs, "baryMat", &baryMat );
//	cvReleaseFileStorage( &fs );
//	printf("Error 8\n");

	// Normalization...
	weight=cvSum( &baryMat ).val[0];	
	
	// ... and barycenter:
	if( weight>1e02 && *convMaxOUT > 1.5e06 ){ /*Weight only equals 0 when we can't find the ball*/
//		printf("Error 9\n");
//		printf("Size of baryMat: %d, %d. Size of xLeverArm: %d, %d\n", baryMat.rows, baryMat.cols, xLeverArm->rows, xLeverArm->cols );
		cvMatMul( &baryMat, xLeverArm, xSumTmp );
		cvMatMul( yLeverArm, &baryMat, ySumTmp );

		// Set coordinates relative to absolute origin.
		cvmSet( y_measrd, 0, 0, cvSum( xSumTmp ).val[0]/weight+(double)cropRect.x ); //cvmSet is a tricky function. It references from 0.
		cvmSet( y_measrd, 1, 0, cvSum( ySumTmp ).val[0]/weight+(double)cropRect.y ); //cvmSet is a tricky function. It references from 0.
//		printf("x=%f, y=%f absolute\n", cvmGet( y_measrd, 0, 0 ), cvmGet( y_measrd, 1, 0 ));
	}
	else 
		std::cout << "Warning, have lost ball " << std::endl;
	
	//Local garbage collection
	cvReleaseMat( &xSumTmp );
	cvReleaseMat( &ySumTmp );
	cvReleaseMat( &xLeverArm );
	cvReleaseMat( &yLeverArm );
	
	//Outputs
	*convMatOUT=cvCloneMat(convMat);
	*bwImageOUT= cvCloneImage(bwImage);
	*edgeImageOUT= cvCloneImage(edgeImage);
	
	// Garbage Collection
	cvResetImageROI( camFrame );
//	free( ballPoint );
	cvReleaseImage( &cropImage );
	cvReleaseImage( &grayImage );
	cvReleaseImage( &bwImage );
	cvReleaseImage( &edgeImage );
	cvReleaseMat( &convMat );

	return;
}

//============================================================================================

void conv2D( IplImage *im1, IplImage *im2, int dft_M, int dft_N )
{
	// This function takes two images and computes the convolution. It was
	// copied directly off of the OpenCV wiki. I made only minor modifications
	// to adjust for variable types.
	
	// Declare structures
	IplImage *edgeImage, *maskImage;
	CvMat *conv, *dft_eIm, *dft_mIm;
	CvMat tmp; //Temporary matrix which is NECESSARY. Won't work without it.
	
	// Create image structures
	edgeImage= cvCreateImage(cvGetSize(im1),  IPL_DEPTH_32F, 1); //Edge matrix in F32
	maskImage = cvCreateImage(cvGetSize(im2),  IPL_DEPTH_32F, 1); //Masque matrix in F32
	conv = cvCreateMat( edgeImage->height + maskImage->height - 1, edgeImage->width + maskImage->width - 1, CV_32F );
//	printf("Error 3\n");

	cvConvert( im1, edgeImage ); //Change from integer to float
	cvConvert( im2, maskImage ); //Change from integer to float
	
	dft_eIm = cvCreateMat( dft_M, dft_N, conv->type );
	dft_mIm = cvCreateMat( dft_M, dft_N, conv->type );
	
	// Copy edgeImage to dft_eIm and pad dft_eIm with zeros
	cvGetSubRect( dft_eIm, &tmp, cvRect(0,0,edgeImage->width,edgeImage->height));
	cvCopy( edgeImage, &tmp, NULL );
	cvGetSubRect( dft_eIm, &tmp, cvRect(edgeImage->width,0,dft_eIm->cols - edgeImage->width,edgeImage->height));  
	cvZero( &tmp );
	
	// Repeat the same with the second array
	cvGetSubRect( dft_mIm, &tmp, cvRect(0,0,maskImage->width,maskImage->height));
	cvCopy( maskImage, &tmp, NULL );
	cvGetSubRect( dft_mIm, &tmp, cvRect(maskImage->width,0,dft_mIm->cols - maskImage->width,maskImage->height));
	cvZero( &tmp );
//	printf("Error 4\n");

	// No need to pad bottom part of dft_eIm or dft_mIm with zeros because of
	// use of non-zero rows parameter in cvDFT() calls below
	cvDFT( dft_eIm, dft_eIm, CV_DXT_FORWARD, edgeImage->height );
	cvDFT( dft_mIm, dft_mIm, CV_DXT_FORWARD, maskImage->height );
	
	//Perform the actual convolution
	cvMulSpectrums( dft_eIm, dft_mIm, dft_eIm, 0 );
	
	//Inverses the DFT into a regular image
	cvDFT( dft_eIm, dft_eIm, CV_DXT_INV_SCALE, conv->rows ); // calculate only the top part
	
	//Get only the region that has data. This should return a matrix of the same size and
	//position as the input image.
	cvGetSubRect( dft_eIm, &tmp, cvRect(maskImage->width/2,maskImage->height/2,edgeImage->width,edgeImage->height) );//Might need fixing.
//	printf("Error 5\n");
	
	//NECESSARY matrix clone. Won't work without it.
	convMat=cvCloneMat(&tmp);
	
	//Garbage Collection
	cvReleaseImage(&edgeImage);
	cvReleaseImage(&maskImage);
	cvReleaseMat(&dft_eIm);
	cvReleaseMat(&dft_mIm);	
	cvReleaseMat(&conv);

}